package com.tfswx.jkgl.core.utils.sm4;

import java.io.OutputStream;

public class EncryptSm4OutputStream extends DecorateOutputStreamBase {

    private final static int SM4_ITEM_LENGTH = 16;
    private final static InnerSM4_Context SM4_CONTEXT_ENCRYPT = new InnerSM4_Context(InnerSM4Stream.SM4_ENCRYPT);

    public EncryptSm4OutputStream(OutputStream out) {
        super(out, 16, 16);
    }

    private final byte[] iv = new byte[16];
    private final byte[] out = new byte[16];

    @Override
    protected void innerOpen() {
        System.arraycopy(SM4_CONTEXT_ENCRYPT.ivBytes, 0, iv, 0, SM4_ITEM_LENGTH);
    }

    @Override
    protected int innerWriteFrom(byte[] sourceBuffer, int sourceLength, byte[] destBuffer) {

        if (sourceLength == SM4_ITEM_LENGTH) {
            // 直接加密
            for (int i = 0; i < SM4_ITEM_LENGTH; i++) {
                out[i] = ((byte) (sourceBuffer[i] ^ iv[i]));
            }
            InnerSM4Stream.sm4_one_round(SM4_CONTEXT_ENCRYPT.sk, out, destBuffer);
            System.arraycopy(destBuffer, 0, iv, 0, SM4_ITEM_LENGTH);
        } else {
            // 最后一段，补位
            InnerSM4Stream.paddingByteArray16(sourceBuffer, sourceLength);
            // 加密
            for (int i = 0; i < SM4_ITEM_LENGTH; i++) {
                out[i] = ((byte) (sourceBuffer[i] ^ iv[i]));
            }
            InnerSM4Stream.sm4_one_round(SM4_CONTEXT_ENCRYPT.sk, out, destBuffer);
        }

        return SM4_ITEM_LENGTH;
    }
}
